import pdb

import torch
import torch.nn as nn
from torch.nn import Parameter
import torch.nn.functional as F
import numpy


class Aggregator(nn.Module):
    def __init__(self, batch_size, dim, dropout, act, name=None):
        super(Aggregator, self).__init__()
        self.dropout = dropout
        self.act = act
        self.batch_size = batch_size
        self.dim = dim

    def forward(self):
        pass


# 从model.py中知道这里 forward 函数的输入分别是会话中所有item的embedding：hidden， item的邻居item矩阵：adj，掩码：mask_item
class LocalAggregator(nn.Module):
    def __init__(self, dim, alpha, dropout=0., name=None):
        super(LocalAggregator, self).__init__()
        self.dim = dim
        self.dropout = dropout

        # 四种边对应的关系向量
        self.a_0 = nn.Parameter(torch.Tensor(self.dim, 1))
        self.a_1 = nn.Parameter(torch.Tensor(self.dim, 1))
        self.a_2 = nn.Parameter(torch.Tensor(self.dim, 1))
        self.a_3 = nn.Parameter(torch.Tensor(self.dim, 1))
        # 偏置项
        self.bias = nn.Parameter(torch.Tensor(self.dim))

        self.leakyrelu = nn.LeakyReLU(alpha)

    def forward(self, hidden, adj, mask_item=None):
        h = hidden  # batch_size*seq_size**hidden_size
        batch_size = h.shape[0]
        # N是seq_len
        N = h.shape[1]

        # 这里是公式(7)中点积部分：h_{v_i} dot h_{v_j}
        # Pytorch中的repeat函数是将tensor对应的维度扩大指定倍数，这里是第0、1、2维分别扩大1、1、N倍变为batch_size*seq_size*(N*hidden_size)
        # Pytorch 中 * 表示矩阵对应位置相乘；这一步的计算可以尝试一个简单的tensor模拟。辅助理解。所有hi与hj两两相乘得到一个batch_size*N*N*self.dim的结果
        a_input = (h.repeat(1, 1, N).view(batch_size, N * N, self.dim) * h.repeat(1, N, 1)) \
            .view(batch_size, N, N, self.dim)

        # 对a_input进行四种不同的线性映射表示四种类型的边，对应公式（7）的a_rij_T(h_{v_i} dot h_{v_j})
        e_0 = torch.matmul(a_input, self.a_0)
        e_1 = torch.matmul(a_input, self.a_1)
        e_2 = torch.matmul(a_input, self.a_2)
        e_3 = torch.matmul(a_input, self.a_3)

        # 进行LeakyReLU激活，完成公式(7)其余部分
        e_0 = self.leakyrelu(e_0).squeeze(-1).view(batch_size, N, N)
        e_1 = self.leakyrelu(e_1).squeeze(-1).view(batch_size, N, N)
        e_2 = self.leakyrelu(e_2).squeeze(-1).view(batch_size, N, N)
        e_3 = self.leakyrelu(e_3).squeeze(-1).view(batch_size, N, N)

        # 创建一个值全为-9e15的矩阵，如果adj中的值为1，则让alpha中对应的值为e_0中对应的值，否则alpha中对应的值为-9e15
        mask = -9e15 * torch.ones_like(e_0)
        alpha = torch.where(adj.eq(1), e_0, mask)
        # 这样alpha即由1和-9e15两个值组成，若其中的值等于2，则alpha中对应的值为e_1中对应的值，否则保持不变，3和4同理
        alpha = torch.where(adj.eq(2), e_1, alpha)
        alpha = torch.where(adj.eq(3), e_2, alpha)
        alpha = torch.where(adj.eq(4), e_3, alpha)
        # 对alpha进行softmax操作，得到注意力分数
        alpha = torch.softmax(alpha, dim=-1)

        # 公式(9)加权求和，得到会话级项目表示
        output = torch.matmul(alpha, h)
        return output


class GlobalAggregator(nn.Module):
    def __init__(self, dim, dropout, act=torch.relu, name=None):
        super(GlobalAggregator, self).__init__()
        self.dropout = dropout
        self.act = act
        self.dim = dim

        # 对应公式2中的W1，这里+1是因为公式2中拼接了一个Wij
        self.w_1 = nn.Parameter(torch.Tensor(self.dim + 1, self.dim))
        # 对应公式2中的q_1^T
        self.w_2 = nn.Parameter(torch.Tensor(self.dim, 1))
        self.w_3 = nn.Parameter(torch.Tensor(2 * self.dim, self.dim))
        self.bias = nn.Parameter(torch.Tensor(self.dim))

    def forward(self, self_vectors, neighbor_vector, batch_size, masks, neighbor_weight, extra_vector=None):
        """
        计算用户的全局偏好，并输出对应的embedding
        :param self_vectors: 输入会话序列的初始embedding
        :param neighbor_vector: 输入会话序列的邻居结点对应的embedding，shape为：batch_size*seq_len*sample_num*self.dim
        :param batch_size: 训练时的batch_size，默认为100
        :param masks: 会话序列掩码，但是没有用到
        :param neighbor_weight: 和邻居结点之间边的权重，shape为batch_size*seq_len*sample_num
        :param extra_vector: 传shape为batch_size*seq_len*self.dim的会话序列均值，对应公式2中的s
        :return: 节点的全局级表示
        """
        # ======== 断点调试 =======
        # pdb.set_trace()
        # ======== 断点调试 =======
        if extra_vector is not None:
            # 这里是公式（2），让s \cdot h_{v_j}，之后再与边权w_{ij}做拼接，之后通过一个线性映射得到alpha=\pi(v_i,v_j)
            alpha = torch.matmul(
                # torch.cat(..., -1)  最后一个维度拼接，拼接之后大小为batch_size*seq_len*sample_num*(hidden_size+1)
                torch.cat([extra_vector.unsqueeze(2).repeat(1, 1, neighbor_vector.shape[2], 1) * neighbor_vector,
                           # 公式2计算结束之后，使用squeeze(-1)将alpha维度降为batch_size*seq_len*sample_num*hidden_size
                           neighbor_weight.unsqueeze(-1)], -1), self.w_1).squeeze(-1)
            # 公式2的激活部分
            alpha = F.leaky_relu(alpha, negative_slope=0.2)
            # 公式2的q_1^T进行线性映射部分
            alpha = torch.matmul(alpha, self.w_2).squeeze(-1)
            # 公式4进行权重归一化
            alpha = torch.softmax(alpha, -1).unsqueeze(-1)
            # 公式1进行加权求和
            neighbor_vector = torch.sum(alpha * neighbor_vector, dim=-2)
        else:
            # 均值池化
            neighbor_vector = torch.mean(neighbor_vector, dim=2)
        # self_vectors = F.dropout(self_vectors, 0.5, training=self.training)
        # 公式5中的拼接部分，聚合得到的邻居信息与更新前的item向量进行拼接，激活再线性映射，从而得到outout
        output = torch.cat([self_vectors, neighbor_vector], -1)
        output = F.dropout(output, self.dropout, training=self.training)
        # 线性变换
        output = torch.matmul(output, self.w_3)
        # ======== 断点调试 =======
        # pdb.set_trace()
        # ======== 断点调试 =======
        # 重新调整维度为batch_size*seq_len*hidden_size
        output = output.view(batch_size, -1, self.dim)
        # 公式5的激活部分
        output = self.act(output)
        # 返回节item的全局级表示，大小为batch_size*seq_len*hidden_size
        return output
